# coding=utf-8
import math
import numpy as np
import tensorflow as tf
slim = tf.contrib.slim
import tf_extended as tfe # tf_extended
import a11_net_params as netparam
import a13_gen_anchor_box

"""
Get the detected bounding boxes from the SSD network output.
"""
def detected_bboxes(predictions, 
                    localisations,
                    select_threshold = None, 
                    nms_threshold = 0.5,
                    clipping_bbox = None, 
                    top_k = 400, 
                    keep_top_k = 200):
    # Select top_k bboxes from predictions, and clip
    rscores, rbboxes = ssd_common.tf_ssd_bboxes_select(
    	predictions, 
    	localisations,
        select_threshold = select_threshold,
        num_classes = self.params.num_classes)
    # sort
    rscores, rbboxes = tfe.bboxes_sort(rscores, rbboxes, top_k=top_k)
    # Apply NMS algorithm.
    rscores, rbboxes = \
        tfe.bboxes_nms_batch(rscores, rbboxes,
                             nms_threshold = nms_threshold,
                             keep_top_k = keep_top_k)
    if clipping_bbox is not None:
        rbboxes = tfe.bboxes_clip(clipping_bbox, rbboxes)
    return rscores, rbboxes

"""
Extract classes, scores and bounding boxes from network output layers.
Batch-compatible: inputs are supposed to have batch-type shapes.

Args:
  predictions_net: List of SSD prediction layers;
  localizations_net: List of localization layers;
  select_threshold: Classification threshold for selecting a box. All boxes
    under the threshold are set to 'zero'. If None, no threshold applied.
Return:
  d_scores, d_bboxes: Dictionary of scores and bboxes Tensors of
    size Batches X N x 1 | 4. Each key corresponding to a class.
"""
def tf_ssd_bboxes_select(predictions_net, localizations_net,
                         select_threshold=None,
                         num_classes=21,
                         ignore_class=0,
                         scope=None):

    with tf.name_scope(scope, 'ssd_bboxes_select',
                       [predictions_net, localizations_net]):
        l_scores = []
        l_bboxes = []
        for i in range(len(predictions_net)):
            scores, bboxes = tf_ssd_bboxes_select_layer(predictions_net[i],
                                                        localizations_net[i],
                                                        select_threshold,
                                                        num_classes,
                                                        ignore_class)
            l_scores.append(scores)
            l_bboxes.append(bboxes)
        # Concat results.
        d_scores = {}
        d_bboxes = {}
        for c in l_scores[0].keys():
            ls = [s[c] for s in l_scores]
            lb = [b[c] for b in l_bboxes]
            d_scores[c] = tf.concat(ls, axis=1)
            d_bboxes[c] = tf.concat(lb, axis=1)
        return d_scores, d_bboxes

"""
Extract classes, scores and bounding boxes from features in one layer.
Batch-compatible: inputs are supposed to have batch-type shapes.

Args:
  predictions_layer: A SSD prediction layer;
  localizations_layer: A SSD localization layer;
  select_threshold: Classification threshold for selecting a box. All boxes
    under the threshold are set to 'zero'. If None, no threshold applied.
Return:
  d_scores, d_bboxes: Dictionary of scores and bboxes Tensors of
    size Batches X N x 1 | 4. Each key corresponding to a class.
"""
def tf_ssd_bboxes_select_layer(predictions_layer, localizations_layer,
                               select_threshold=None,
                               num_classes=21,
                               ignore_class=0,
                               scope=None):

    select_threshold = 0.0 if select_threshold is None else select_threshold
    with tf.name_scope(scope, 'ssd_bboxes_select_layer',
                       [predictions_layer, localizations_layer]):
        # Reshape features: Batches x N x N_labels | 4
        p_shape = tfe.get_shape(predictions_layer)
        predictions_layer = tf.reshape(predictions_layer,
                                       tf.stack([p_shape[0], -1, p_shape[-1]]))
        l_shape = tfe.get_shape(localizations_layer)
        localizations_layer = tf.reshape(localizations_layer,
                                         tf.stack([l_shape[0], -1, l_shape[-1]]))

        d_scores = {}
        d_bboxes = {}
        for c in range(0, num_classes):
            if c != ignore_class:
                # Remove boxes under the threshold.
                scores = predictions_layer[:, :, c]
                fmask = tf.cast(tf.greater_equal(scores, select_threshold), scores.dtype)
                scores = scores * fmask
                bboxes = localizations_layer * tf.expand_dims(fmask, axis=-1)
                # Append to dictionary.
                d_scores[c] = scores
                d_bboxes[c] = bboxes

        return d_scores, d_bboxes

if __name__ == '__main__':
	
	print 'OK'






